import unittest

from question_classifier import *
from question_parser import *
from answer_search import *


class TestModule(unittest.TestCase):
    # 实体链接模块 单元测试
    def test_entity_linker(self):
        classifier = QuestionClassifier()
        question = "毛囊炎病人可以吃什么食物？"
        res_classify = classifier.classify(question)
        linked_entities = res_classify['args']
        self.assertEqual("毛囊炎", linked_entities)

    # 关系预测模块 单元测试
    def test_relation_classifier(self):
        classifier = QuestionClassifier()
        question = "色汗症病人是去皮肤科挂号吗？"
        res_classify = classifier.classify(question)
        relation = res_classify['question_types']
        self.assertEqual("is_belonging", relation)

    # 逻辑形式构造模块 单元测试
    def test_parser(self):
        parser = QuestionPaser()
        res_classify = {'args': {'肾虚': ['disease']}, 'question_types': ['disease_cause', 'disease_prevent']}
        result = parser.parser_main(res_classify)
        res_sql = [
            {'question_type': 'disease_cause', 'sql': ["MATCH (m:Disease) where m.name = '肾虚' return m.name, m.cause"]},
            {'question_type': 'disease_prevent',
             'sql': ["MATCH (m:Disease) where m.name = '肾虚' return m.name, m.prevent"]}]
        self.assertEqual(res_sql[0]['question_type'], result[0]['question_type'])
        self.assertEqual(res_sql[0]['sql'], result[0]['sql'])
        self.assertEqual(res_sql[1]['question_type'], result[1]['question_type'])
        self.assertEqual(res_sql[1]['sql'], result[1]['sql'])

    # 自然语言答案构造模块 单元测试
    def test_searcher(self):
        searcher = AnswerSearcher()
        res_sql = [{
            'question_type': 'disease_cureway',
            'entities': ['肾虚'],
            'sql': ["MATCH (m:Disease) where m.name = '肾虚' return m.name, m.cure_way"]
        }]
        answer = ["肾虚可以尝试如下治疗：['中医治疗', '支持性治疗']。\n"]
        final_answers = searcher.search_main(res_sql)
        self.assertEqual(answer, final_answers['answer'])
